package com.junxonline.instant.common.sharding;

import cn.hutool.core.date.DateUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.google.common.collect.Lists;
import com.junxonline.instant.common.constant.RedisConstant;
import com.junxonline.instant.common.util.common.RedisUtils;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingValue;
import org.springframework.core.env.Environment;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.*;

/**
 * 按月范围分片算法
 * 命名规则：表名_yyyyMM
 *
 * @author JunX
 * @date 2021-10-27
 */
public class MonthRangeShardingAlgorithm implements RangeShardingAlgorithm<Date> {

    private RedisUtils redisUtils;

    @Override
    public Collection<String> doSharding(Collection<String> collection, RangeShardingValue<Date> rangeShardingValue) {

        if (redisUtils == null) {
            redisUtils = SpringUtil.getBean(RedisUtils.class);
        }

        Set<String> tableNames = new HashSet<>(collection);
        String tableName = Lists.newArrayList(collection).get(0);
        if (rangeShardingValue != null) {
            Date startDate = rangeShardingValue.getValueRange().lowerEndpoint();
            Date endDate = rangeShardingValue.getValueRange().upperEndpoint();
            if (endDate.after(new Date())) {
                endDate = DateUtil.date();
            }
            List<String> names = new ArrayList<>();
            // 先查询缓存中是否存在
            if (redisUtils.hasKey(RedisConstant.KEY_DATABASE_TABLES)) {
                List<Object> cacheTables = redisUtils.lGet(RedisConstant.KEY_DATABASE_TABLES, 0, -1);
                if (cacheTables != null && cacheTables.size() > 0) {
                    for (Object table : cacheTables) {
                        names.add((String) table);
                    }
                }
            }
            // 不存在则查询数据库
            else {
                names = this.getAllTableNameBySchema();
            }
            for (; startDate.before(endDate) || startDate.equals(endDate); startDate = DateUtil.tomorrow()) {
                String key = DateUtil.format(startDate, "yyyyMM");
                String searchTableName = tableName.concat("_").concat(key);
                for (String name : names) {
                    if (searchTableName.equals(name)) {
                        tableNames.add(name);
                    }
                }
            }
        }
        return tableNames;
    }

    private List<String> getAllTableNameBySchema() {
        List<String> tableNames = new ArrayList<>();
        try {
            Environment environment = SpringUtil.getApplicationContext().getEnvironment();
            Connection connection = DriverManager.getConnection(
                    Objects.requireNonNull(environment.getProperty("spring.shardingsphere.datasource.ms0.url")),
                    Objects.requireNonNull(environment.getProperty("spring.shardingsphere.datasource.ms0.username")),
                    Objects.requireNonNull(environment.getProperty("spring.shardingsphere.datasource.ms0.password"))
            );
            Statement statement = connection.createStatement();
            ResultSet resultSet = statement.executeQuery("SHOW TABLES;");
            while (resultSet.next()) {
                tableNames.add(resultSet.getString(1));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return tableNames;
    }

}
